#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Excel工具类
提供Excel文件的读取、写入、格式设置等功能
"""

import os
import sys
from openpyxl import load_workbook, Workbook
from openpyxl.styles import Font, Alignment, Border, Side, PatternFill
from openpyxl.utils import get_column_letter
import pandas as pd
import numpy as np
from datetime import datetime

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.config.constants import EXPORT_SETTINGS
from src.config.exceptions import ExcelProcessingError


class ExcelUtils:
    """Excel工具类，提供Excel文件的读取、写入、格式设置等功能"""

    def __init__(self):
        """初始化Excel工具类"""
        # 设置默认字体
        self.default_font = Font(name="Arial", size=11)

        # 设置默认对齐方式
        self.default_alignment = Alignment(
            horizontal="left", vertical="center", wrap_text=True
        )

        # 设置默认边框
        self.default_border = Border(
            left=Side(style="thin"),
            right=Side(style="thin"),
            top=Side(style="thin"),
            bottom=Side(style="thin"),
        )

        # 设置默认填充
        self.default_fill = PatternFill(fill_type=None)

    def read_excel(self, file_path, sheet_name=0, header=0, usecols=None, dtype=None):
        """读取Excel文件数据

        Args:
            file_path: Excel文件路径
            sheet_name: 工作表名称或索引，默认为0（第一个工作表）
            header: 表头行索引，默认为0
            usecols: 要读取的列，默认为None（读取所有列）
            dtype: 指定列的数据类型，默认为None

        Returns:
            DataFrame: 读取到的数据

        Raises:
            ExcelProcessingError: Excel处理错误
        """
        try:
            # 检查文件是否存在
            if not os.path.exists(file_path):
                raise ExcelProcessingError(f"文件不存在: {file_path}")

            # 使用pandas读取Excel文件
            df = pd.read_excel(
                file_path,
                sheet_name=sheet_name,
                header=header,
                usecols=usecols,
                dtype=dtype,
            )

            return df

        except Exception as e:
            raise ExcelProcessingError(f"读取Excel文件失败: {str(e)}")

    def read_excel_with_merged_cells_info(self, file_path, sheet_name=0):
        """读取Excel文件并获取合并单元格信息

        Args:
            file_path: Excel文件路径
            sheet_name: 工作表名称或索引，默认为0（第一个工作表）

        Returns:
            tuple: (DataFrame, dict)，包含数据和合并单元格信息

        Raises:
            ExcelProcessingError: Excel处理错误
        """
        try:
            # 检查文件是否存在
            if not os.path.exists(file_path):
                raise ExcelProcessingError(f"文件不存在: {file_path}")

            # 使用openpyxl读取工作簿
            wb = load_workbook(file_path, data_only=True)

            # 获取指定工作表
            if isinstance(sheet_name, int):
                sheet = wb.worksheets[sheet_name]
            else:
                sheet = wb[sheet_name]

            # 读取表头和数据行
            data = []
            headers = []

            # 获取表头
            for cell in sheet[1]:
                if cell.value is not None:
                    headers.append(cell.value)
                else:
                    headers.append("")

            # 获取数据行
            for row in sheet.iter_rows(min_row=2):
                row_data = []
                for cell in row:
                    row_data.append(cell.value)
                data.append(row_data)

            # 使用pandas创建DataFrame
            df = pd.DataFrame(data, columns=headers)

            # 获取合并单元格信息
            merged_cells_info = {}

            for merged_cell_range in sheet.merged_cells.ranges:
                # 获取合并单元格的边界
                min_row, min_col, max_row, max_col = merged_cell_range.bounds

                # 转换为0-based索引
                df_min_row = min_row - 2  # 减去表头行
                df_max_row = max_row - 2  # 减去表头行
                df_min_col = min_col - 1  # 转换为0-based列索引
                df_max_col = max_col - 1  # 转换为0-based列索引

                # 检查行索引是否在DataFrame范围内
                if df_min_row >= 0 and df_max_row < len(df):
                    # 检查列索引是否有效
                    if df_min_col >= 0 and df_max_col < len(df.columns):
                        # 获取合并单元格的左上角单元格值
                        top_left_value = df.iloc[df_min_row, df_min_col]

                        # 构建合并单元格信息
                        column_name = df.columns[df_min_col]

                        if column_name not in merged_cells_info:
                            merged_cells_info[column_name] = []

                        merged_cells_info[column_name].append(
                            {
                                "start_row": df_min_row,
                                "end_row": df_max_row,
                                "value": top_left_value,
                            }
                        )

            return df, merged_cells_info

        except Exception as e:
            raise ExcelProcessingError(
                f"读取Excel文件并获取合并单元格信息失败: {str(e)}"
            )

    def export_to_excel(
        self,
        df,
        file_path,
        sheet_name="Sheet1",
        index=False,
        merge_cells_info=None,
    ):
        """导出数据到Excel文件

        Args:
            df: 要导出的数据
            file_path: 输出文件路径
            sheet_name: 工作表名称，默认为'Sheet1'
            index: 是否包含索引列，默认为False
            merge_cells_info: 合并单元格信息，默认为None

        Raises:
            ExcelProcessingError: Excel处理错误
        """
        try:
            # 确保目录存在
            os.makedirs(os.path.dirname(file_path), exist_ok=True)

            # 创建ExcelWriter对象
            with pd.ExcelWriter(file_path, engine="openpyxl") as writer:
                # 导出数据
                df.to_excel(writer, sheet_name=sheet_name, index=index)

                # 获取工作表对象
                sheet = writer.sheets[sheet_name]

                # 设置列宽
                for i, col in enumerate(df.columns):
                    # 计算列宽
                    column_width = max(
                        df[col].astype(str).map(len).max(), len(str(col)) + 2
                    )
                    # 设置列宽，最大不超过50
                    column_width = min(column_width, 50)
                    # 设置列宽
                    sheet.column_dimensions[get_column_letter(i + 1)].width = (
                        column_width
                    )

                # 设置表头样式
                for cell in sheet[1]:
                    cell.font = Font(bold=True)
                    cell.alignment = Alignment(
                        horizontal="center", vertical="center"
                    )
                    cell.border = self.default_border

                # 应用默认样式到所有数据单元格
                for row in sheet.iter_rows(min_row=2, max_row=len(df) + 1):
                    for cell in row:
                        cell.font = self.default_font
                        cell.alignment = self.default_alignment
                        cell.border = self.default_border

                # 如果有规格型号列，设置为文本格式
                if "规格型号" in df.columns:
                    # 获取规格型号列的索引
                    spec_col_index = df.columns.get_loc("规格型号") + 1
                    # 设置整列为文本格式
                    for row in sheet.iter_rows(
                        min_row=2,
                        max_row=len(df) + 1,
                        min_col=spec_col_index,
                        max_col=spec_col_index,
                    ):
                        for cell in row:
                            if isinstance(
                                cell.value, (int, float)
                            ) and not pd.isna(cell.value):
                                cell.number_format = "@"
                                if cell.value == int(cell.value):
                                    cell.value = str(int(cell.value))
                                else:
                                    cell.value = str(cell.value)
        except Exception as e:
            raise ExcelProcessingError(f"导出Excel文件失败: {str(e)}")

    def create_workbook(self):
        """创建一个新的工作簿

        Returns:
            Workbook: 新创建的工作簿对象
        """
        return Workbook()

    def save_workbook(self, workbook, file_path):
        """保存工作簿到文件

        Args:
            workbook: 工作簿对象
            file_path: 文件路径

        Raises:
            ExcelProcessingError: Excel处理错误
        """
        try:
            # 确保目录存在
            os.makedirs(os.path.dirname(file_path), exist_ok=True)

            # 保存工作簿
            workbook.save(file_path)

        except Exception as e:
            raise ExcelProcessingError(f"保存工作簿失败: {str(e)}")

    def create_worksheet(self, workbook, sheet_name):
        """在工作簿中创建一个新的工作表

        Args:
            workbook: 工作簿对象
            sheet_name: 工作表名称

        Returns:
            Worksheet: 新创建的工作表对象
        """
        return workbook.create_sheet(title=sheet_name)

    def merge_cells(self, worksheet, start_row, start_col, end_row, end_col):
        """合并单元格

        Args:
            worksheet: 工作表对象
            start_row: 开始行
            start_col: 开始列
            end_row: 结束行
            end_col: 结束列
        """
        worksheet.merge_cells(
            start_row=start_row,
            start_column=start_col,
            end_row=end_row,
            end_column=end_col,
        )

    def set_cell_value(self, worksheet, row, col, value):
        """设置单元格的值

        Args:
            worksheet: 工作表对象
            row: 行号
            col: 列号
            value: 要设置的值
        """
        worksheet.cell(row=row, column=col).value = value

    def get_cell_value(self, worksheet, row, col):
        """获取单元格的值

        Args:
            worksheet: 工作表对象
            row: 行号
            col: 列号

        Returns:
            单元格的值
        """
        return worksheet.cell(row=row, column=col).value

    def apply_format_to_cell(
        self,
        worksheet,
        row,
        col,
        font=None,
        alignment=None,
        border=None,
        fill=None,
    ):
        """应用格式到单元格

        Args:
            worksheet: 工作表对象
            row: 行号
            col: 列号
            font: 字体，默认为None
            alignment: 对齐方式，默认为None
            border: 边框，默认为None
            fill: 填充，默认为None
        """
        cell = worksheet.cell(row=row, column=col)

        if font is not None:
            cell.font = font

        if alignment is not None:
            cell.alignment = alignment

        if border is not None:
            cell.border = border

        if fill is not None:
            cell.fill = fill

    def apply_format_to_range(
        self,
        worksheet,
        start_row,
        start_col,
        end_row,
        end_col,
        font=None,
        alignment=None,
        border=None,
        fill=None,
    ):
        """应用格式到单元格范围

        Args:
            worksheet: 工作表对象
            start_row: 开始行
            start_col: 开始列
            end_row: 结束行
            end_col: 结束列
            font: 字体，默认为None
            alignment: 对齐方式，默认为None
            border: 边框，默认为None
            fill: 填充，默认为None
        """
        # 对单元格范围内的每个单元格应用格式
        for row in range(start_row, end_row + 1):
            for col in range(start_col, end_col + 1):
                self.apply_format_to_cell(
                    worksheet,
                    row,
                    col,
                    font=font,
                    alignment=alignment,
                    border=border,
                    fill=fill,
                )

    def format_excel_data(self, data, columns=None):
        """格式化Excel数据

        Args:
            data: 要格式化的数据
            columns: 列名列表，默认为None

        Returns:
            DataFrame: 格式化后的数据
        """
        # 如果数据不是DataFrame，尝试转换为DataFrame
        if not isinstance(data, pd.DataFrame):
            try:
                df = pd.DataFrame(data, columns=columns)
            except Exception as e:
                raise ExcelProcessingError(f"格式化数据失败: {str(e)}")
        else:
            df = data.copy()

        # 处理日期列
        for col in df.columns:
            if pd.api.types.is_datetime64_any_dtype(df[col]):
                df[col] = df[col].dt.strftime("%Y-%m-%d")

        # 处理数字列
        numeric_columns = df.select_dtypes(include="number").columns
        for col in numeric_columns:
            # 检查是否所有值都是整数
            if np.allclose(df[col].dropna(), df[col].dropna().astype(int)):
                df[col] = df[col].astype(pd.Int64Dtype())

        return df
